#include "TriangleMesh.cuh"

using namespace OBR;

static __forceinline__ __device__ float3 calcBaryCoord(float3 a, float3 b, float3 c, float3 p,
                                                       float3 n)
{
    float3 bary;
    float areaABC = dot(n, cross((b - a), (c - a)));
    float areaPBC = dot(n, cross((b - p), (c - p)));
    float areaPCA = dot(n, cross((c - p), (a - p)));
    bary.x = areaPBC / areaABC;      // alpha
    bary.y = areaPCA / areaABC;      // beta
    bary.z = 1.0f - bary.x - bary.y; // gamma
    return bary;
}

extern "C" __global__ void __closesthit__triangle_mesh()
{
    uint p0, p1;
    SamplePRD *prd = getPRDWithPayload(p0, p1);
    prd->hit = true;

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("==== triangle hit program launched\n");
    }
#endif

    // SBT info
    HitData *sbt_record = reinterpret_cast<HitData *>(optixGetSbtDataPointer());
    const uint mesh_id = sbt_record->mesh_id;
    const uint mesh_instance_id = sbt_record->mesh_instance_id;

    // trace info
    uint tri_idx = optixGetPrimitiveIndex();
    float3 ray_dir = optixGetWorldRayDirection();
    float3 ray_orig = optixGetWorldRayOrigin();
    float ray_t = optixGetRayTmax();
    float3 hit_pos = ray_orig + ray_t * ray_dir;
    float2 bary2 = optixGetTriangleBarycentrics();

    // init surface interaction
    InteractionData &si = prd->inter;
    si.type = InteractionType::SURFACE_INTERACTION;
    si.wo = prd->inter.wi;
    si.wo.t = ray_t;
    si.mat_id = sbt_record->mat_id;
    si.mat_instance_id = sbt_record->mat_instance_id;
    si.light_id = sbt_record->light_id;
    si.light_instance_id = sbt_record->light_instance_id;
    si.med_id_inside = sbt_record->med_id_inside;
    si.med_instance_id_inside = sbt_record->med_instance_id_inside;
    si.med_id_outside = sbt_record->med_id_outside;
    si.med_instance_id_outside = sbt_record->med_instance_id_outside;
    si.world_pos = hit_pos;
    si.mesh_id = mesh_id;
    si.mesh_instance_id = mesh_instance_id;

    /* ------------------- calculate surface interaction -------------------- */
    TriangleMeshData &meshdata =
        reinterpret_cast<TriangleMeshData *>(sbt_record->hit_payload)[mesh_instance_id];

    const Triangle &tri = meshdata.triangles[tri_idx];
    const Vertex &v0 = meshdata.vertexes[tri.vertexes[0]];
    const Vertex &v1 = meshdata.vertexes[tri.vertexes[1]];
    const Vertex &v2 = meshdata.vertexes[tri.vertexes[2]];

    float3 bary = make_float3(1.0f - bary2.x - bary2.y, bary2.x, bary2.y);
    float3 n_obj = normalize(bary.x * v0.n + bary.y * v1.n + bary.z * v2.n);
    si.n = normalize(optixTransformNormalFromObjectToWorldSpace(n_obj));
    si.uv = bary.x * v0.tex_uv + bary.y * v1.tex_uv + bary.z * v2.tex_uv;
    si.obj_pos = optixTransformPointFromWorldToObjectSpace(hit_pos);

    if (prd->ray_occlusion_test)
    {
        // in occlusion test we dont need differentials
        return;
    }

    /* ----------------------- calculate differentials ---------------------- */
    if (si.wo.has_differential)
    {
        // compute normal differentials
        float2 duv02 = v0.tex_uv - v2.tex_uv;
        float2 duv12 = v1.tex_uv - v2.tex_uv;
        float3 dn02 = v0.n - v2.n;
        float3 dn12 = v1.n - v2.n;
        // NOTE: can we precompute dpdu / dpdv ?
        float3 dp02 = v0.pos - v2.pos;
        float3 dp12 = v1.pos - v2.pos;
        float determinant = duv02.x * duv12.y - duv02.y * duv12.x;
        if (fabs(determinant) < 1e-8)
        {
            si.dpdu = si.dpdv = si.dndu = si.dndv = make_float3(0.0f);
        }
        else
        {
            float inv_det = 1.0f / determinant;
            si.dndu = (duv12.y * dn02 - duv02.y * dn12) * inv_det;
            si.dndv = (duv02.x * dn12 - duv12.x * dn02) * inv_det;
            si.dpdu = (duv12.y * dp02 - duv02.y * dp12) * inv_det;
            si.dpdv = (duv02.x * dp12 - duv12.x * dp02) * inv_det;
        }

        // Estimate screen space change in $\pt{}$ and $(u,v)$
        // Compute auxiliary intersection points with plane
        float3 p = si.world_pos;
        float d = dot(si.n, p);
        float tx = -(dot(si.n, si.wo.rx_o) - d) / dot(si.n, si.wo.rx_d);
        if (isinf(tx) || isnan(tx))
            goto fail;
        float ty = -(dot(si.n, si.wo.ry_o) - d) / dot(si.n, si.wo.ry_d);
        if (isinf(ty) || isnan(ty))
            goto fail;
        float3 px = si.wo.rx_o + si.wo.rx_d * tx;
        float3 py = si.wo.ry_o + si.wo.ry_d * ty;
        si.dpdx = px - p;
        si.dpdy = py - p;

        // Compute $(u,v)$ offsets at auxiliary points
        // Choose two dimensions to use for ray offset computation
        int dim[2];
        if (fabs(si.n.x) > fabs(si.n.y) && fabs(si.n.x) > fabs(si.n.z))
        {
            dim[0] = 1;
            dim[1] = 2;
        }
        else if (fabs(si.n.y) > fabs(si.n.z))
        {
            dim[0] = 0;
            dim[1] = 2;
        }
        else
        {
            dim[0] = 0;
            dim[1] = 1;
        }

        // Initialize _A_, _Bx_, and _By_ matrices for offset computation
        float A[2][2] = {get_float3_at(&(si.dpdu), dim[0]), get_float3_at(&(si.dpdv), dim[0]),
                         get_float3_at(&(si.dpdu), dim[1]), get_float3_at(&(si.dpdv), dim[1])};
        float Bx[2] = {get_float3_at(&(si.dpdx), dim[0]), get_float3_at(&(si.dpdx), dim[1])};
        float By[2] = {get_float3_at(&(si.dpdy), dim[0]), get_float3_at(&(si.dpdy), dim[1])};
        if (!solve_linear_system_2x2(A, Bx, &si.dudx, &si.dvdx))
            si.dudx = si.dvdx = 0;
        if (!solve_linear_system_2x2(A, By, &si.dudy, &si.dvdy))
            si.dudy = si.dvdy = 0;
    }
    else
    {
    fail:
        si.dndu = si.dndv = si.dpdu = si.dpdv = si.dpdx = si.dpdy = make_float3(0.0f);
        si.dudx = si.dudy = si.dvdx = si.dvdy = 0.0f;
    }

    return;
}